using System.Collections.Generic;
using UnityEngine;
using System.IO;
using System.Text;
using UnityEditor;
using UnityEditor.Rendering;

namespace Mars.ShaderAnalysis
{
    public class ShaderAnalysis
    {
        public bool isExpanded = false;
        public Shader shader;
        public List<SubShaderAnalysis> subShaderAnalysisList = new List<SubShaderAnalysis>();
        public string shaderPath;
        public string shaderFullPath;
        public List<ShaderKeywordsState> keywordsStateList = new List<ShaderKeywordsState>();
        private string countString = "";
        private string combinationsString = "";

        private static StringBuilder stringBuilder = new StringBuilder();
        private static StringBuilder debugStringBuilder = new StringBuilder();

        /// <summary>
        /// 设置Shader
        /// </summary>
        /// <param name="shader"></param>
        public void SetShader(Shader shader)
        {
            this.shader = shader;
            shaderPath = AssetDatabase.GetAssetPath(shader);
            shaderFullPath = EditorTools.projectPath + shaderPath;
            shaderFullPath = shaderFullPath.Replace("/", "\\");
        }

        /// <summary>
        /// 重置
        /// </summary>
        public void Reset()
        {
            keywordsStateList.Clear();
            countString = "";
            combinationsString = "";
        }

        /// <summary>
        /// 获取对应Keyword状态
        /// </summary>
        public ShaderKeywordsState GetState(string keyword)
        {
            for (int i = 0; i < keywordsStateList.Count; i++)
            {
                ShaderKeywordsState state = keywordsStateList[i];
                if (state.keywordString == keyword)
                {
                    return state;
                }
            }

            return null;
        }

        /// <summary>
        /// 获取数值字符串
        /// </summary>
        public string GetCountString()
        {
            if (shader == null)
            {
                return "";
            }

            if (!string.IsNullOrEmpty(countString))
            {
                return countString;
            }

            int totalCount = keywordsStateList.Count;
            int useCount = 0;
            for (int i = 0; i < keywordsStateList.Count; i++)
            {
                ShaderKeywordsState state = keywordsStateList[i];
                if (state != null && state.isOpen)
                {
                    useCount++;
                }
            }

            stringBuilder.Clear();
            stringBuilder.Append(totalCount.ToString("00"));
            stringBuilder.Append(" [<color=#00ff00>");
            stringBuilder.Append(useCount.ToString("00"));
            stringBuilder.Append("</color>/<color=#ff0000>");
            stringBuilder.Append((totalCount - useCount).ToString("00"));
            stringBuilder.Append("</color>]");
            countString = stringBuilder.ToString();
            return countString;
        }

        /// <summary>
        /// 获取变体字符串
        /// </summary>
        public string GetCombinationsString()
        {
            if (shader == null)
            {
                return "";
            }

            if (!string.IsNullOrEmpty(combinationsString))
            {
                return combinationsString;
            }

            Debug.Log("分析Shader : " + shader.name);
            Debug.Log("顶点组合数计算");
            ulong totalVertexCombineCount = 0;
            //顶点着色器变体组合数
            for (int i = 0; i < subShaderAnalysisList.Count; i++)
            {
                SubShaderAnalysis subShaderAnalysis = subShaderAnalysisList[i];
                Debug.Log("    SubShader_" + (i + 1));
                List<PassAnalysis> passAnalysisList = subShaderAnalysis.passAnalysisList;
                for (int j = 0; j < passAnalysisList.Count; j++)
                {
                    PassAnalysis passAnalysis = passAnalysisList[j];
                    ulong combineCount = 1;
                    int count = 0;
                    debugStringBuilder.Clear();
                    debugStringBuilder.Append("\n");
                    for (int k = 0; k < passAnalysis.keywordsStateList.Count; k++)
                    {
                        ShaderKeywordsState state = passAnalysis.keywordsStateList[k];
                        if (state.isOpen && state.hasVertex)
                        {
                            combineCount *= (ulong)state.keywords.Count;
                            count++;
                            debugStringBuilder.Append("                ");
                            debugStringBuilder.Append(state.keywordString);
                            debugStringBuilder.Append(" : ");
                            debugStringBuilder.Append(state.keywords.Count);
                            debugStringBuilder.Append("\n");
                        }
                    }

                    Debug.Log("        " + passAnalysis.passName + " (" + passAnalysis.lightMode + ") = " + combineCount + " , 复杂度 " + count + debugStringBuilder);
                    totalVertexCombineCount += combineCount;
                }
            }

            stringBuilder.Clear();
            stringBuilder.Append("顶点组合数 = ");
            stringBuilder.Append(EditorGUILayoutUtility.LargeNumberToString(totalVertexCombineCount));
            stringBuilder.Append("  ;  ");

            Debug.Log("片元组合数计算");
            ulong totalFragmentCombineCount = 0;
            //顶点着色器变体组合数
            for (int i = 0; i < subShaderAnalysisList.Count; i++)
            {
                SubShaderAnalysis subShaderAnalysis = subShaderAnalysisList[i];
                Debug.Log("    SubShader_" + (i + 1));
                List<PassAnalysis> passAnalysisList = subShaderAnalysis.passAnalysisList;
                for (int j = 0; j < passAnalysisList.Count; j++)
                {
                    PassAnalysis passAnalysis = passAnalysisList[j];
                    ulong combineCount = 1;
                    int count = 0;
                    debugStringBuilder.Clear();
                    debugStringBuilder.Append("\n");
                    for (int k = 0; k < passAnalysis.keywordsStateList.Count; k++)
                    {
                        ShaderKeywordsState state = passAnalysis.keywordsStateList[k];
                        if (state.isOpen && state.hasFragment)
                        {
                            combineCount *= (ulong)state.keywords.Count;
                            count++;
                            debugStringBuilder.Append("                ");
                            debugStringBuilder.Append(state.keywordString);
                            debugStringBuilder.Append(" : ");
                            debugStringBuilder.Append(state.keywords.Count);
                            debugStringBuilder.Append("\n");
                        }
                    }

                    Debug.Log("        " + passAnalysis.passName + " (" + passAnalysis.lightMode + ") = " + combineCount + " , 复杂度 " + count + debugStringBuilder);
                    totalFragmentCombineCount += combineCount;
                }
            }

            stringBuilder.Append("片元组合数 = ");
            stringBuilder.Append(EditorGUILayoutUtility.LargeNumberToString(totalFragmentCombineCount));
            stringBuilder.Append("  ;  ");

            int multiCompileCount = 0;
            for (int i = 0; i < keywordsStateList.Count; i++)
            {
                ShaderKeywordsState state = keywordsStateList[i];
                if (state.isOpen && state.isMultiCompile)
                {
                    multiCompileCount++;
                }
            }

            stringBuilder.Append("MultiCompile数 = ");
            stringBuilder.Append(multiCompileCount);

            combinationsString = stringBuilder.ToString();
            return combinationsString;
        }

        /// <summary>
        /// 分析Keywords
        /// </summary>
        public void Execute()
        {
            if (shader == null)
            {
                return;
            }

            keywordsStateList.Clear();
            subShaderAnalysisList.Clear();
            int subShaderIndex = -1;
            int passIndex = -1;
            HashSet<string> keywordsHashSet = new HashSet<string>();
            string[] shaderTexts = File.ReadAllLines(shaderFullPath);
            Dictionary<int, HashSet<string>> outOfPassKeywords = new Dictionary<int, HashSet<string>>();
            for (int i = 0; i < shaderTexts.Length; i++)
            {
                string lineText = shaderTexts[i].Trim();
                //去掉两个编辑器的Keyword
                if (lineText.Contains("EDITOR") || lineText.Contains("DEBUG"))
                {
                    continue;
                }

                //统计LightMode
                if (lineText.ToLower().StartsWith("subshader"))
                {
                    subShaderIndex++;
                    passIndex = -1;
                    subShaderAnalysisList.Add(new SubShaderAnalysis());
                }

                if (lineText.ToLower().Equals("pass"))
                {
                    passIndex++;
                    subShaderAnalysisList[subShaderIndex].passAnalysisList.Add(new PassAnalysis());
                }

                if (lineText.StartsWith("Name \""))
                {
                    string[] splits = lineText.Split("\"");
                    string passName = splits[1];
                    if (subShaderIndex < 0 || passIndex < 0)
                    {
                        Debug.LogError(shader.name + " 分析失败，" + lineText + " 需要写在Pass里才能分析成功。");
                        continue;
                    }

                    subShaderAnalysisList[subShaderIndex].passAnalysisList[passIndex].passName = passName;
                }

                if (lineText.Contains("\"LightMode\""))
                {
                    string[] splits = lineText.Split("\"");
                    string lightMode = splits[3];
                    if (subShaderIndex < 0 || passIndex < 0)
                    {
                        Debug.LogError(shader.name + " 分析失败，" + lineText + " 需要写在Pass里才能分析成功。");
                        continue;
                    }

                    subShaderAnalysisList[subShaderIndex].passAnalysisList[passIndex].lightMode = lightMode;
                }

                //统计keyword
                if (lineText.StartsWith("#pragma shader_feature") || lineText.StartsWith("#pragma multi_compile"))
                {
                    if (passIndex < 0)
                    {
                        if (!outOfPassKeywords.ContainsKey(subShaderIndex))
                        {
                            outOfPassKeywords.Add(subShaderIndex, new HashSet<string>());
                        }

                        outOfPassKeywords[subShaderIndex].Add(lineText);
                        continue;
                    }

                    ShaderKeywordsState keywordsState = new ShaderKeywordsState(lineText, true);
                    keywordsState.Analysis();
                    if (!keywordsHashSet.Contains(lineText))
                    {
                        keywordsStateList.Add(keywordsState);
                        keywordsHashSet.Add(lineText);
                    }

                    subShaderAnalysisList[subShaderIndex].passAnalysisList[passIndex].keywordsStateList.Add(keywordsState);
                }

                if (lineText.StartsWith("//#pragma shader_feature") || lineText.StartsWith("//#pragma multi_compile"))
                {
                    if (passIndex < 0)
                    {
                        if (!outOfPassKeywords.ContainsKey(subShaderIndex))
                        {
                            outOfPassKeywords.Add(subShaderIndex, new HashSet<string>());
                        }

                        outOfPassKeywords[subShaderIndex].Add(lineText);
                        continue;
                    }

                    if (subShaderIndex >= 0 && passIndex < 0)
                    {
                        Debug.LogError(shader.name + " 分析失败，Keyword字符串 " + lineText + " 需要写在Pass里 或 写在所有SubShader前 才能分析成功。");
                        continue;
                    }

                    lineText = lineText.Replace("//", string.Empty);
                    ShaderKeywordsState keywordsState = new ShaderKeywordsState(lineText, false);
                    keywordsState.Analysis();
                    if (!keywordsHashSet.Contains(lineText))
                    {
                        keywordsStateList.Add(keywordsState);
                        keywordsHashSet.Add(lineText);
                    }

                    subShaderAnalysisList[subShaderIndex].passAnalysisList[passIndex].keywordsStateList.Add(keywordsState);
                }
            }

            //写在所有Pass前的，所有pass都添加
            foreach (var kv in outOfPassKeywords)
            {
                int index = kv.Key;
                foreach (var s in kv.Value)
                {
                    string lineText = s;
                    bool isOpen = true;
                    if (lineText.StartsWith("//#pragma shader_feature") || lineText.StartsWith("//#pragma multi_compile"))
                    {
                        lineText = lineText.Replace("//", string.Empty);
                        isOpen = false;
                    }

                    ShaderKeywordsState keywordsState = new ShaderKeywordsState(lineText, isOpen);
                    keywordsState.Analysis();
                    if (!keywordsHashSet.Contains(lineText))
                    {
                        keywordsStateList.Add(keywordsState);
                        keywordsHashSet.Add(lineText);
                    }

                    for (int i = 0; i < subShaderAnalysisList.Count; i++)
                    {
                        if (index == -1 || index == i)
                        {
                            for (int j = 0; j < subShaderAnalysisList[i].passAnalysisList.Count; j++)
                            {
                                //为了防止引用同一个对象，这里新建一个新的
                                ShaderKeywordsState state = new ShaderKeywordsState(lineText, isOpen);
                                state.Analysis();
                                subShaderAnalysisList[i].passAnalysisList[j].keywordsStateList.Add(state);
                                // subShaderAnalysisList[i].passAnalysisList[j].keywordsStateList.Add(keywordsState);
                            }
                        }
                    }
                }
            }

            //排序
            keywordsStateList.Sort((a, b) =>
            {
                if (a.isCustom && !b.isCustom)
                {
                    return 1;
                }

                if (!a.isCustom && b.isCustom)
                {
                    return -1;
                }

                return a.keywordString.CompareTo(b.keywordString);
            });
        }

        /// <summary>
        /// 开关Keyword
        /// </summary>
        public void SetKeywordActive(string keyword, bool isActive)
        {
            if (shader == null || string.IsNullOrEmpty(shaderFullPath))
            {
                return;
            }

            for (int i = 0; i < keywordsStateList.Count; i++)
            {
                if (keywordsStateList[i].keywordString == keyword)
                {
                    keywordsStateList[i].isOpen = isActive;
                    string shaderAllText = File.ReadAllText(shaderFullPath);
                    if (isActive)
                    {
                        shaderAllText = shaderAllText.Replace("//" + keyword, keyword);
                    }
                    else
                    {
                        shaderAllText = shaderAllText.Replace(keyword, "//" + keyword);
                    }

                    //写入文本
                    File.WriteAllText(shaderFullPath, shaderAllText);
                    //刷新资源
                    EditorUtility.SetDirty(shader);
                    AssetDatabase.SaveAssets();
                    AssetDatabase.Refresh();
                    //重新分析
                    Reset();
                    Execute();
                    return;
                }
            }
        }

        /// <summary>
        /// 编译并显示Pass
        /// </summary>
        public void CompileVariantAndAnalysisPass(ShaderCompilerPlatform shaderCompilerPlatform, BuildTarget buildTarget, bool isDisplayDialog = false)
        {
            for (int i = 0; i < subShaderAnalysisList.Count; i++)
            {
                List<PassAnalysis> passAnalysisList = subShaderAnalysisList[i].passAnalysisList;
                for (int j = 0; j < passAnalysisList.Count; j++)
                {
                    PassAnalysis passAnalysis = passAnalysisList[j];
                    if (isDisplayDialog)
                    {
                        EditorUtility.DisplayProgressBar("编译并分析 " + (j + 1) + "/" + passAnalysisList.Count, shader.name + " : " + passAnalysis.passName + " (" + passAnalysis.lightMode + ")", (float)j / passAnalysisList.Count);
                    }

                    passAnalysis.CompileVariantAndAnalysis(shader, i, j, shaderCompilerPlatform, buildTarget);
                }
            }

            EditorUtility.ClearProgressBar();
        }
    }
}